import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
sns.set(style="whitegrid")

# 1. 读取数据
data = pd.read_csv('../data/train.csv')

print("===== 数据基本信息 =====")
print(f"数据形状：{data.shape}")
print(data.info())

# 2. 自动识别目标列 Attrition
target_col = None
for col in data.columns:
    if col.lower() == 'attrition':
        target_col = col
        break

if target_col:
    print(f"\n🎯 检测到目标列: {target_col}")
    print("目标列分布：")
    print(data[target_col].value_counts(normalize=True))

    plt.figure(figsize=(4,4))
    sns.countplot(x=target_col, data=data, palette='Set2')
    plt.title(f'{target_col} 分布')
    plt.show()
else:
    print("⚠️ 未检测到 Attrition 列，请确认数据集目标列名。")

# 3. 特征类型识别
numeric_features = data.select_dtypes(include=['int64', 'float64']).columns.tolist()
categorical_features = data.select_dtypes(include=['object']).columns.tolist()

# 排除目标列
if target_col in numeric_features:
    numeric_features.remove(target_col)
if target_col in categorical_features:
    categorical_features.remove(target_col)

# 4. 进一步识别“有序分类”特征
# 这里使用启发式规则：如果唯一值数量较少且是数值类型，则视为有序变量
ordinal_features = []
for col in numeric_features:
    unique_vals = data[col].nunique()
    if 3 <= unique_vals <= 10 and data[col].dtype in ['int64', 'float64']:
        ordinal_features.append(col)

# 去掉有序列，剩下才是纯数值型
pure_numeric_features = [c for c in numeric_features if c not in ordinal_features]

# 5. 输出分类结果
print("\n===== 自动特征分类结果 =====")
print(f"数值型特征（Numeric）：{pure_numeric_features}")
print(f"有序型特征（Ordinal）：{ordinal_features}")
print(f"无序分类特征（Nominal）：{categorical_features}")

# 6. 类别型特征样本预览
print("\n===== 无序分类特征示例 =====")
for col in categorical_features:
    print(f"{col} : {data[col].unique()[:5]}")

# 7. 数值特征相关性热力图
if len(pure_numeric_features) > 1:
    plt.figure(figsize=(12,8))
    sns.heatmap(data[pure_numeric_features].corr(), cmap='coolwarm', annot=False)
    plt.title('数值特征相关性热力图')
    plt.show()

# 8. 输出总结
print("\n===== 总结 =====")
print(f"共检测到 {len(pure_numeric_features)} 个数值特征, {len(ordinal_features)} 个有序特征, {len(categorical_features)} 个无序特征。")
if target_col:
    print(f"目标列为: {target_col}")
